#!/bin/bash
set -euo pipefail

# 脚本名称：check_fstype.sh

# 功能说明：
# 本脚本用于检查虚拟磁盘或物理磁盘分区的文件系统类型。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 blkid、lsblk、qemu-nbd 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 检查文件系统类型。
# 5. 清理挂载的 nbd 设备（如果有的话）。

# 注意事项：
# - 需要 root 权限运行部分命令，如 blkid、lsblk、qemu-nbd。
# - 在执行操作前，建议备份重要数据，以防万一。

# 帮助信息
usage() {
  echo "用途: 检查磁盘或分区的文件系统类型"
  echo "用法: $(basename ${BASH_SOURCE}) <设备>"
  echo "示例: $(basename ${BASH_SOURCE}) /dev/sdb"
  echo "      $(basename ${BASH_SOURCE}) /dev/sdb1"
  exit 1
}

#--------- 依赖管理 ---------#
install_dependencies() {
    local required=(blkid lsblk qemu-nbd partprobe)
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -gt 0 ]; then
        echo "缺少依赖: ${missing[*]}，正在安装..."
        sudo apt-get update -qq
        sudo apt-get install -y -qq qemu-utils util-linux >/dev/null
    fi
}

#--------- 全局变量 ---------#
nbd_device=""
virtual_disk_attached=false

# 清理函数：断开当前连接的 NBD 设备
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "断开设备: $nbd_device"
        sudo qemu-nbd -d "$nbd_device" >/dev/null
        virtual_disk_attached=false
    fi
}
trap cleanup EXIT SIGINT SIGTERM

#--------- NBD设备管理 ---------#
# 从 /dev/nbd* 中选择一个空闲的设备
select_nbd() {
    for dev in /dev/nbd*; do
        # 检查设备是否存在对应的 pid 文件且内容为空（表示未被占用）
        pid_file="/sys/block/$(basename "$dev")/pid"
        if [ -f "$pid_file" ]; then
            if [ ! -s "$pid_file" ]; then
                echo "$dev"
                return 0
            fi
        else
            # 若 pid 文件不存在，则认为设备可用
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘到NBD设备
setup_nbd() {
    local disk_path="$1"
    echo "加载 nbd 模块..."

    # 尝试加载nbd模块
    if ! sudo modprobe nbd max_part=16 2>/dev/null; then
        echo "nbd模块加载失败，尝试安装内核扩展模块..."
        sudo apt-get install -y "linux-modules-extra-$(uname -r)"
        if ! sudo modprobe nbd max_part=16; then
            echo "错误：无法加载nbd模块！"
            echo "请手动执行：sudo apt install linux-modules-extra-$(uname -r)"
            exit 1
        fi
    fi

    nbd_device=$(select_nbd)
    if [ -z "$nbd_device" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi

    echo "连接 $disk_path 到 $nbd_device..."
    sudo qemu-nbd -c "$nbd_device" "$disk_path"
    sleep 2  # 等待设备初始化
    virtual_disk_attached=true
}

# 使用 blkid 检查文件系统类型
check_fstype_by_blkid() {
  echo "方法 1 (blkid):"
  sudo blkid "$1" -s TYPE -o value || echo "使用 blkid 获取文件系统类型失败。"
}

# 使用 lsblk 检查文件系统类型
check_fstype_by_lsblk() {
  echo "方法 2 (lsblk):"
  sudo lsblk -no FSTYPE "$1" || echo "使用 lsblk 获取文件系统类型失败。"
}

# 检查设备上的所有分区
check_all_partitions() {
  local device="$1"
  local partitions=$(lsblk -rno NAME "$device" | grep -E "^$(basename "$device")p?[0-9]+$")
  
  for part in $partitions; do
    local partition="/dev/$part"
    echo "检查分区: $partition"
    check_fstype_by_blkid "$partition"
    check_fstype_by_lsblk "$partition"
  done
}

# 主函数
main() {
  install_dependencies

  # 检查是否提供了设备参数
  if [ $# -eq 0 ]; then
    usage
  fi

  # 设备参数
  local device="$1"

  # 检查提供的设备是否存在
  if [ -b "$device" ]; then
    echo "选择的物理磁盘或分区: $device"
    target_device="$device"
  elif [ -f "$device" ]; then
    echo "选择的虚拟磁盘镜像文件: $device"
    setup_nbd "$device"
    target_device="$nbd_device"
  else
    echo "错误： '$device' 既不是块设备也不是普通文件。"
    exit 1
  fi

  # 检查目标是磁盘还是分区
  if lsblk -no TYPE "$target_device" | grep -q "disk"; then
    echo "检测到磁盘: $target_device"
    # 检查所有分区
    echo "检查 $target_device 上的所有分区"
    check_all_partitions "$target_device"
  else
    echo "检测到分区: $target_device"
    # 检查分区
    echo "运行 blkid 检查 $target_device"
    check_fstype_by_blkid "$target_device"
    echo "运行 lsblk 检查 $target_device"
    check_fstype_by_lsblk "$target_device"
  fi
}

# 检查脚本是否被直接执行（而不是被 source）
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  main "$@"
else
  export -f check_fstype_by_blkid check_fstype_by_lsblk
fi
